package cn.yhp.chatgpt.data.trigger.http;

import cn.yhp.chatgpt.data.domain.auth.service.IAuthService;
import cn.yhp.chatgpt.data.domain.openai.model.aggregates.ChatProcessAggregate;
import cn.yhp.chatgpt.data.domain.openai.service.IChatService;
import cn.yhp.chatgpt.data.trigger.http.dto.ChatGPTRequestDTO;
import cn.yhp.chatgpt.data.types.common.Constants;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * @author yhp
 * @description
 * @date 2024/2/1
 */
@RestController
@Slf4j
@RequestMapping("/api/${app.config.api-version}/chatglm")
public class ChatGPTAIServiceController {

    @Resource
    private IChatService chatService;
    @Resource
    private IAuthService authService;

    @Resource
    private ThreadPoolExecutor threadPoolExecutor;

    /**
     * 流式回答问题
     * curl -X POST \
     * http://localhost:8090/api/v1/chatglm/chat/completions \
     * -H 'Content-Type: application/json;charset=utf-8'
     * -H 'Authorization: uuid' \
     * -d '{"messages": [{"content": "写一个java二分查找","role": "user"}],"model": "chatglm-turbo"}'
     * @param dto
     * @param token
     * @param response
     * @return
     */
    @PostMapping("chat/completions")
    public ResponseBodyEmitter completionsStream(@RequestBody ChatGPTRequestDTO dto,
                                                 @RequestHeader("Authorization") String token,
                                                 HttpServletResponse response){
        // 1.基础配置：流式输出，编码，禁用缓存
        response.setContentType("text/event-stream");
        response.setCharacterEncoding("utf-8");
        response.setHeader("Cache-Control","no-cache");

        // 2、构建异步响应对象(默认超时3min)
        ResponseBodyEmitter emitter = new ResponseBodyEmitter(3 * 60 * 1000L);
        // 3. 对 Token 过期拦截
        boolean success = authService.checkToken(token);

        if (!success) {
            try {
                emitter.send(Constants.ResponseCode.TOKEN_ERROR.getInfo());
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            emitter.complete();
            return emitter;
        }
        // 3. 获取 OpenID
        String openid = authService.openid(token);
        log.info("流式问答请求处理，openid:{} 请求模型:{}", openid, dto.getModel());

        // 4.构建参数
        ChatProcessAggregate processAggregate = ChatProcessAggregate.builder()
                .openid(openid)
                .model(dto.getModel())
                .messages(dto.getMessages())
                .build();

        // 3.请求
        return chatService.completions(processAggregate,emitter);
    }

    @RequestMapping(value = "/chat", method = RequestMethod.GET)
    public ResponseBodyEmitter completionsStream(HttpServletResponse response) {
        response.setContentType("text/event-stream");
        response.setCharacterEncoding("UTF-8");
        response.setHeader("Cache-Control", "no-cache");

        ResponseBodyEmitter emitter = new ResponseBodyEmitter();

        threadPoolExecutor.execute(() -> {
            for (int i = 0; i < 10; i++) {
                try {
                    emitter.send("strdddddddddddddddd\r\n" + i);
                    Thread.sleep(800);
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
            }
            emitter.complete();
        });

        return emitter;
    }
}
